#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
比亚迪股票数据分析
此脚本用于分析比亚迪后复权历史行情数据，完成以下任务：
1. 读取数据并进行基本分析
2. 计算相关系数并绘制热力图
3. 绘制2023年全年收盘价时序图
4. 绘制2023年4-6月的K线图
"""

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import mplfinance as mpf
from datetime import datetime

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 定义文件路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_FILE = os.path.join(BASE_DIR, '比亚迪后复权历史行情数据.xlsx')

class BYDStockAnalyzer:
    """
    比亚迪股票数据分析类
    用于分析比亚迪后复权历史行情数据
    """
    
    def __init__(self):
        """初始化分析器"""
        self.data = None
        self.processed_data = None
    
    def load_data(self):
        """
        加载Excel数据
        """
        print("=== 开始加载比亚迪后复权历史行情数据 ===")
        try:
            self.data = pd.read_excel(DATA_FILE)
            print(f"数据加载成功！数据形状：{self.data.shape} 行 × {self.data.shape[1]} 列")
            print("\n数据前5行：")
            print(self.data.head())
            return True
        except Exception as e:
            print(f"数据加载失败：{e}")
            return False
    
    def preprocess_data(self):
        """
        预处理数据
        """
        if self.data is None:
            print("错误：请先加载数据")
            return False
        
        # 深拷贝原始数据
        self.processed_data = self.data.copy()
        
        # 假设日期列名为'日期'，如果不是可以修改
        date_column = None
        for col in self.processed_data.columns:
            if '日期' in col or 'date' in col.lower():
                date_column = col
                break
        
        if date_column:
            # 转换日期列
            try:
                self.processed_data[date_column] = pd.to_datetime(self.processed_data[date_column])
                self.processed_data = self.processed_data.set_index(date_column)
                print(f"已将'{date_column}'设置为索引并转换为日期格式")
            except Exception as e:
                print(f"日期转换失败：{e}")
        
        return True
    
    def descriptive_analysis(self):
        """
        对定量变量进行描述性分析
        """
        if self.processed_data is None:
            print("错误：请先预处理数据")
            return False
        
        print("\n=== 定量变量描述性分析 ===")
        numeric_cols = self.processed_data.select_dtypes(include=['float64', 'int64']).columns
        
        if len(numeric_cols) > 0:
            description = self.processed_data[numeric_cols].describe()
            print(description)
            return True
        else:
            print("未找到定量变量")
            return False
    
    def correlation_analysis(self):
        """
        计算变量之间的相关系数，并绘制热力图
        """
        if self.processed_data is None:
            print("错误：请先预处理数据")
            return False
        
        print("\n=== 变量相关系数分析 ===")
        numeric_cols = self.processed_data.select_dtypes(include=['float64', 'int64']).columns
        
        if len(numeric_cols) < 2:
            print("定量变量不足，无法计算相关系数")
            return False
        
        # 计算相关系数
        corr_matrix = self.processed_data[numeric_cols].corr()
        print("相关系数矩阵：")
        print(corr_matrix)
        
        # 绘制热力图
        plt.figure(figsize=(12, 10))
        sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt='.2f', linewidths=0.5)
        plt.title('变量相关系数热力图')
        plt.tight_layout()
        
        # 保存图像
        heatmap_file = os.path.join(BASE_DIR, '相关系数热力图.png')
        plt.savefig(heatmap_file, dpi=300)
        print(f"相关系数热力图已保存至：{heatmap_file}")
        plt.close()
        
        return True
    
    def plot_2023_close_price(self):
        """
        绘制2023年全年的收盘价时序图
        """
        if self.processed_data is None:
            print("错误：请先预处理数据")
            return False
        
        print("\n=== 绘制2023年收盘价时序图 ===")
        
        # 筛选2023年的数据
        try:
            # 假设索引是日期
            if isinstance(self.processed_data.index, pd.DatetimeIndex):
                data_2023 = self.processed_data[(self.processed_data.index.year == 2023)]
            else:
                # 尝试寻找日期列
                date_column = None
                for col in self.processed_data.columns:
                    if pd.api.types.is_datetime64_any_dtype(self.processed_data[col]):
                        date_column = col
                        break
                
                if date_column:
                    data_2023 = self.processed_data[
                        (self.processed_data[date_column].dt.year == 2023)
                    ]
                else:
                    print("未找到日期列，无法筛选2023年数据")
                    return False
            
            if data_2023.empty:
                print("没有找到2023年的数据")
                return False
            
            # 寻找收盘价列
            close_column = None
            for col in data_2023.columns:
                if '收盘' in col or 'close' in col.lower():
                    close_column = col
                    break
            
            if close_column:
                # 绘制时序图
                plt.figure(figsize=(15, 8))
                plt.plot(data_2023.index, data_2023[close_column], 'b-', linewidth=1.5)
                plt.title('2023年比亚迪股票收盘价时序图')
                plt.xlabel('日期')
                plt.ylabel(f'{close_column}')
                plt.grid(True, alpha=0.3)
                plt.tight_layout()
                
                # 保存图像
                time_series_file = os.path.join(BASE_DIR, '2023年收盘价时序图.png')
                plt.savefig(time_series_file, dpi=300)
                print(f"2023年收盘价时序图已保存至：{time_series_file}")
                plt.close()
                return True
            else:
                print("未找到收盘价列")
                return False
        except Exception as e:
            print(f"绘制时序图失败：{e}")
            return False
    
    def plot_k_line_chart(self):
        """
        绘制2023年4月到2023年6月的K线图
        """
        if self.processed_data is None:
            print("错误：请先预处理数据")
            return False
        
        print("\n=== 绘制2023年4-6月K线图 ===")
        
        try:
            # 筛选2023年4-6月的数据
            if isinstance(self.processed_data.index, pd.DatetimeIndex):
                data_k = self.processed_data[
                    (self.processed_data.index.year == 2023) & 
                    (self.processed_data.index.month >= 4) & 
                    (self.processed_data.index.month <= 6)
                ]
            else:
                # 尝试寻找日期列
                date_column = None
                for col in self.processed_data.columns:
                    if pd.api.types.is_datetime64_any_dtype(self.processed_data[col]):
                        date_column = col
                        break
                
                if date_column:
                    data_k = self.processed_data[
                        (self.processed_data[date_column].dt.year == 2023) &
                        (self.processed_data[date_column].dt.month >= 4) &
                        (self.processed_data[date_column].dt.month <= 6)
                    ]
                    # 设置日期为索引
                    data_k = data_k.set_index(date_column)
                else:
                    print("未找到日期列，无法筛选数据")
                    return False
            
            if data_k.empty:
                print("没有找到2023年4-6月的数据")
                return False
            
            # 寻找OHLC列（开盘价、最高价、最低价、收盘价）
            ohlc_cols = {'Open': None, 'High': None, 'Low': None, 'Close': None}
            
            for col in data_k.columns:
                col_lower = col.lower()
                if 'open' in col_lower or '开盘' in col:
                    ohlc_cols['Open'] = col
                elif 'high' in col_lower or '最高' in col:
                    ohlc_cols['High'] = col
                elif 'low' in col_lower or '最低' in col:
                    ohlc_cols['Low'] = col
                elif 'close' in col_lower or '收盘' in col:
                    ohlc_cols['Close'] = col
            
            # 检查是否找到所有必要的列
            if None in ohlc_cols.values():
                missing = [k for k, v in ohlc_cols.items() if v is None]
                print(f"未找到以下必要列：{missing}")
                return False
            
            # 重命名列以便mplfinance使用
            k_data = data_k.rename(columns={
                ohlc_cols['Open']: 'Open',
                ohlc_cols['High']: 'High',
                ohlc_cols['Low']: 'Low',
                ohlc_cols['Close']: 'Close'
            })
            
            # 检查是否有成交量列
            volume_column = None
            for col in data_k.columns:
                if 'volume' in col.lower() or '成交量' in col or '成交' in col:
                    volume_column = col
                    break
            
            # 如果有成交量列，添加到数据中
            if volume_column:
                k_data = k_data.rename(columns={volume_column: 'Volume'})
                # 绘制带成交量的K线图
                mpf.plot(
                    k_data, 
                    type='candle', 
                    volume=True, 
                    title='2023年4-6月比亚迪股票K线图',
                    ylabel='价格',
                    ylabel_lower='成交量',
                    figratio=(15, 8),
                    tight_layout=True,
                    style='yahoo',
                    savefig=dict(fname=os.path.join(BASE_DIR, '2023年4-6月K线图.png'), dpi=300)
                )
            else:
                # 绘制不带成交量的K线图
                mpf.plot(
                    k_data[['Open', 'High', 'Low', 'Close']], 
                    type='candle', 
                    title='2023年4-6月比亚迪股票K线图',
                    ylabel='价格',
                    figratio=(15, 8),
                    tight_layout=True,
                    style='yahoo',
                    savefig=dict(fname=os.path.join(BASE_DIR, '2023年4-6月K线图.png'), dpi=300)
                )
                print("警告：未找到成交量列")
            
            print("2023年4-6月K线图已保存至：", os.path.join(BASE_DIR, '2023年4-6月K线图.png'))
            return True
        except Exception as e:
            print(f"绘制K线图失败：{e}")
            return False
    
    def analyze_k_line_chart(self):
        """
        分析K线图
        """
        print("\n=== K线图分析 ===")
        print("K线图展示了2023年4月至6月比亚迪股票的价格走势，包括开盘价、最高价、最低价和收盘价。")
        print("\nK线图分析要点：")
        print("1. 观察整体趋势：上升趋势（阳线多于阴线且价格逐渐抬高）、下降趋势（阴线多于阳线且价格逐渐走低）或横盘整理。")
        print("2. 识别关键形态：如锤头线、射击之星、吞没形态等可能预示趋势反转的K线组合。")
        print("3. 关注成交量变化：价格上涨时成交量放大通常表示上涨动能强劲，价格下跌时成交量放大可能表示抛压较重。")
        print("4. 支撑位和阻力位：观察价格在哪些位置反复受到支撑或阻力。")
        
        # 生成K线图分析报告
        analysis_file = os.path.join(BASE_DIR, 'K线图分析报告.txt')
        with open(analysis_file, 'w', encoding='utf-8') as f:
            f.write("比亚迪股票2023年4-6月K线图分析报告\n")
            f.write("========================\n\n")
            f.write("K线图展示了2023年4月至6月比亚迪股票的价格走势，包括开盘价、最高价、最低价和收盘价。\n\n")
            f.write("K线图分析要点：\n")
            f.write("1. 观察整体趋势：上升趋势（阳线多于阴线且价格逐渐抬高）、下降趋势（阴线多于阳线且价格逐渐走低）或横盘整理。\n")
            f.write("2. 识别关键形态：如锤头线、射击之星、吞没形态等可能预示趋势反转的K线组合。\n")
            f.write("3. 关注成交量变化：价格上涨时成交量放大通常表示上涨动能强劲，价格下跌时成交量放大可能表示抛压较重。\n")
            f.write("4. 支撑位和阻力位：观察价格在哪些位置反复受到支撑或阻力。\n")
            f.write("\n分析建议：\n")
            f.write("- 结合基本面分析和技术指标（如MACD、RSI等）进行综合判断。\n")
            f.write("- 关注市场整体环境和行业动态对股票价格的影响。\n")
            f.write("- 制定合理的风险控制策略。\n")
        
        print(f"K线图分析报告已保存至：{analysis_file}")


def main():
    """
    主函数，执行完整的比亚迪股票数据分析流程
    """
    print("=== 开始执行比亚迪股票数据分析任务 ===")
    
    # 创建分析器实例
    analyzer = BYDStockAnalyzer()
    
    try:
        # 任务1：读取数据并进行基本分析
        print("\n=== 任务1：读取数据并进行基本分析 ===")
        if not analyzer.load_data():
            print("任务1失败")
            return
        
        # 预处理数据
        if not analyzer.preprocess_data():
            print("数据预处理失败")
            return
        
        # 描述性分析
        if not analyzer.descriptive_analysis():
            print("描述性分析失败")
        
        # 任务2：计算相关系数并绘制热力图
        print("\n=== 任务2：计算相关系数并绘制热力图 ===")
        if not analyzer.correlation_analysis():
            print("任务2失败")
        
        # 任务3：绘制2023年全年收盘价时序图
        print("\n=== 任务3：绘制2023年全年收盘价时序图 ===")
        if not analyzer.plot_2023_close_price():
            print("任务3失败")
        
        # 任务4：绘制2023年4-6月K线图并分析
        print("\n=== 任务4：绘制2023年4-6月K线图并分析 ===")
        if not analyzer.plot_k_line_chart():
            print("绘制K线图失败")
        else:
            # 分析K线图
            analyzer.analyze_k_line_chart()
        
        print("\n=== 比亚迪股票数据分析任务完成 ===")
        print("所有结果已保存至chapter4文件夹")
        
    except Exception as e:
        print(f"执行过程中发生错误：{e}")


if __name__ == "__main__":
    main()